import os
import chromadb
from typing import Optional
from langchain_chroma import Chroma
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore

from ..llms.emb import create_openai_emb, EmbConfig


class ChromaVectorStore:
    @classmethod
    def build(cls,
              collection_name: str = 'default',
              embedding: Optional[Embeddings | EmbConfig] = None,
              persist_directory: Optional[str] = None,
              host: Optional[str] = None,
              port: Optional[int] = None,
              collection_metadata: dict = {
                  "hnsw:space": "ip",
                  "hnsw:search_ef": 100
              },
              **kwargs) -> Chroma:
        if embedding is None:
            embedding_function = create_openai_emb(EmbConfig.default())
        elif isinstance(embedding, EmbConfig):
            embedding_function = create_openai_emb(embedding)
        elif isinstance(embedding, Embeddings):
            embedding_function = embedding
        else:
            raise ValueError(f"Invalid embedding type: {type(embedding)}")

        client = None
        local_directory = None

        if persist_directory:
            # 优先使用本地目录
            local_directory = persist_directory
        elif host:
            # 如果提供了服务器URL，创建客户端
            client_settings = chromadb.config.Settings()
            if port:
                client = chromadb.HttpClient(host=host, port=port, settings=client_settings)
            else:
                client = chromadb.HttpClient(host=host, settings=client_settings)
        else:
            local_directory = os.path.join(os.getenv("TMP_PATH", "tmp"), "chroma_db")

        vector_store = Chroma(
            collection_name=collection_name,
            embedding_function=embedding_function,
            persist_directory=local_directory,
            client=client,
            collection_metadata=collection_metadata,
            **kwargs
        )

        return vector_store
